/*
    Self Deletion Technique
    Author: @5mukx
*/


#![allow(dead_code)]
use std::ffi::{OsStr, OsString};
use std::os::windows::ffi::OsStrExt;
use std::ptr::null_mut;
use winapi::ctypes::c_void;
use winapi::um::fileapi::{CreateFileW, SetFileInformationByHandle, FILE_RENAME_INFO};
use winapi::um::handleapi::CloseHandle;
use winapi::um::heapapi::HeapFree;
use winapi::um::minwinbase::{FileDispositionInfo, FileRenameInfo};
use winapi::um::winbase::{MoveFileExW, FILE_FLAG_DELETE_ON_CLOSE};
use winapi::um::winnt::{FILE_ATTRIBUTE_NORMAL, HEAP_ZERO_MEMORY};
use winapi::um::{fileapi::FILE_DISPOSITION_INFO, heapapi::{GetProcessHeap, HeapAlloc}};

// One time self delete after completing operations
pub(crate) fn self_destruction(){
    // This stream name is for alternate stream.
    let stream = ":smukx_stream";
    let stream_wide: Vec<u16> = OsString::from(stream).encode_wide().chain(std::iter::once(0)).collect();

    unsafe{
        let mut delete_file = FILE_DISPOSITION_INFO{
            DeleteFile: 1 // default value is 1
        };

        let rename_info_size = std::mem::size_of::<FILE_RENAME_INFO>() + (stream_wide.len() * std::mem::size_of::<u16>());
        
        // allocate memory for FILE_RENAME_INFO
        let rename_info = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, rename_info_size) as *mut FILE_RENAME_INFO;

        if rename_info.is_null() {
            panic!("Memory allocation failed");
        }

        delete_file.DeleteFile = 1;
        (*rename_info).FileNameLength = (stream_wide.len() * std::mem::size_of::<u16>()) as u32 - 2;

        std::ptr::copy_nonoverlapping(
            stream_wide.as_ptr(),
            (*rename_info).FileName.as_mut_ptr(),
            stream_wide.len(),
        );

        let path = std::env::current_exe().unwrap();
        let mut full_path: Vec<u16> = OsString::from(path).encode_wide().collect();
        full_path.push(0);

        // Open the file with delete and synchronize permissions
        let handle = CreateFileW(
            full_path.as_ptr(),
            0x00010000 | 0x00100000,
            0x00000001,
            core::ptr::null_mut(),
            3, // OPEN_EXISTING
            0 as u32, // FILE_ATTRIBUTE_NORMAL
            null_mut(),
        );

        if handle.is_null() {
            panic!("Failed to open file");
        }

        // rename the alternate data stream
        let setfileinfohandle = SetFileInformationByHandle(
            handle, 
            FileRenameInfo, 
            rename_info as *mut c_void, 
            rename_info_size as u32
        );

        if setfileinfohandle == 0 {
            panic!("Failed to set file information for rename");
        }

        CloseHandle(handle);

        // Re-open the file with delete on close flag
        let handle = CreateFileW(
            full_path.as_ptr(),
            0x00010000 | 0x40000000 | 0x00100000,
            0x00000001,
            std::ptr::null_mut(),
            3, 
            FILE_ATTRIBUTE_NORMAL | FILE_FLAG_DELETE_ON_CLOSE,
            null_mut(),
        );

        if handle == null_mut() {
            panic!("Failed to re-open file for deletion");
        }

        // mark the file for deletion
        let setfileinfo = SetFileInformationByHandle(
            handle, 
            FileDispositionInfo, 
            &delete_file as *const FILE_DISPOSITION_INFO as *mut c_void, 
            std::mem::size_of::<FILE_DISPOSITION_INFO>() as u32
        );

        if setfileinfo == 0 {
            panic!("Failed to mark file for deletion");
        }

        CloseHandle(handle); // -> delete the file  

        // free the allocated memory
        HeapFree(
            GetProcessHeap(), 
            0, 
            rename_info as *mut c_void,
        );
    }
}

pub(crate) fn schedule_self_delete() {
    let exe_path = std::env::current_exe().unwrap();
    let exe_path_wide: Vec<u16> = OsStr::new(exe_path.to_str().unwrap())
        .encode_wide()
        .chain(Some(0).into_iter())
        .collect();

    unsafe {
        MoveFileExW(
            exe_path_wide.as_ptr(),
            null_mut(),
            0x00000004,
        );
    }
}

